import matplotlib.pyplot as plt
import matplotlib.animation as animation
import torch
from torch import nn

'''
data prepare
'''
x = torch.arange(0, 100, dtype=torch.float32, step=1)
k = 20
b = 100

noise = torch.randint(0, 10, (100, 1))
x = x.view(-1, 1)
y = k * x + b + noise

'''
model prepare
'''


class model(nn.Module):
    def __init__(self):
        super(model, self).__init__()
        self.fc = nn.Linear(1, 1, True)


    def forward(self, x):
        return self.fc(x)


function = model()

'''
LOSS OPTIMIZER 
'''
MSEloss = nn.MSELoss()
optimizer = torch.optim.Adam(function.parameters(), lr=100)

plt.ion()

for i in range(1000):
    result = function(x)
    optimizer.zero_grad()
    loss = MSEloss(result, y)
    loss.backward()
    optimizer.step()

    # print("loss:",loss.item())
    # print("Weight:", function.fc.weight.item())
    # print("Bias:", function.fc.bias.item())
    plt.cla()
    plt.title("iter " + str(i))
    plt.scatter(x.data, y.data, c="red", label="gt")
    plt.plot(x.data, result.data, c="blue", label="pred")
    plt.legend(loc='lower right', fontsize=10)
    plt.text(1, -0.6, f'y={round(function.fc.weight.item(), 4)}*x+{round(function.fc.bias.item(), 4)}')

    fig=plt.gcf()
    fig.savefig("result/line"+str(i)+".jpg")
    plt.pause(0.5)
plt.ioff()

plt.show()


